nodes: skip self.model.to(device_to) walk in unpatch_model on mmap'd weights - #445
nodes: skip self.model.to(device_to) walk in unpatch_model on mmap'd weights#445Booyaka101 wants to merge 3 commits into
Conversation
…ap walk The override forwards unpatch_weights=True to comfy core's unpatch_model, which then calls self.model.to(device_to). Walking that .to() over the GGUF-quantized tensors that are still mmap-backed crashes with a Windows access violation (city96#444), surfacing whenever ComfyUI actually invokes unpatch_model — easiest repro is wrapping UnetLoaderGGUF in a node that returns multiple outputs, which flushes cache and triggers the unpatch path. The maintainer mitigated this in Sep 2024 (commit 6dbb4ba) by passing device_to=None to super, then reverted in 717a0e1 because that broke VRAM estimation. Reporter's proposed fix (passing unpatch_weights=False to super) avoids the crash but skips comfy core's non-quantized backup restoration — patch_weight_to_device's `if key not in self.backup` guard then never re-saves the original weight, and the next run's temp_weight = weight.to(...) reads the still-patched value, so LoRA deltas compound across runs. Restore the non-quantized backups in the override mirroring comfy core's loop (copy_to_param / set_attr_param + backup.clear() + comfy_patched_weights cleanup), then forward unpatch_weights=False so super skips its weight-restore + model.to(device_to) block. Keeps the crash fix without regressing LoRA correctness or VRAM accounting. Verified with a stub-based reproducer (verify_unpatch.py) covering three scenarios: pre-fix crashes on simulated mmap .to() — matches city96#444 reporter fix avoids crash, but non-quantized weight stays patched this fix avoids crash AND restores non-quantized backup correctly Closes city96#444 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
Confirmed working on real CUDA. RTX 3060 (12GB), Windows 10, ComfyUI 0.22.2, PyTorch 2.8.0+cu129. Running Flux 2 Klein 9b, with both TE and diffusion model as GGUFs. Was facing an issue where any attempt to unload GGUF model weights (either manually or automatically, in my case triggered by normal between-run model caching) led to the same "Windows fatal exception: access violation" fault described above. After applying this patch, both manual and automatic unloads work and repeated runs are stable. LoRAs produce identical output across multiple runs with the same seed. |
|
After further usage, one edge case to report: changing LoRA strength between runs produces output as if no LoRA is applied at all. Appears the quantized tensors' .patches are cleared correctly on unpatch, but not re-attached on the subsequent in-place re-patch since the model stays resident (unsure about this, though). A full model reload between strength changes fixes it. |
…override The previous override forwarded unpatch_weights=False to super to dodge the self.model.to(device_to) walk that faults on mmap'd quantized tensors, but that also skipped the model_loaded_weight_memory / model_offload_buffer_memory reset super does in the same block. partially_load checks that counter immediately after calling unpatch_model and short-circuits self.load() if it is nonzero, so subsequent LoRA strength changes left .patches cleared without re-attaching anything (no LoRA applied at all on re-runs). Inline the full unpatch_weights body from base ModelPatcher minus that one device walk: hooks/pin/lowvram cleanup, backup restore, uuid reset, memory counters, comfy_patched_weights deletion. Upstream tracked at Comfy-Org/ComfyUI#14142.
Booyaka101
changed the title
|
@romybaby thanks for catching this, the regression is real and caused by the previous version of this PR. What happened: forwarding Just pushed 32414a1 which inlines the whole Worth retesting your strength-change workflow with this commit if you have a moment. |
|
Thanks for the quick response! The commit works perfectly on my end, changing strengths between runs now correctly reapplies the LoRA at the new strength level. |
|
Additional unload-correctness data point for the current PR shape (single host, single A/B run). On ComfyUI 0.33.1 + PyTorch 2.13.0+cu130, I tested a downstream override with the same relevant behavior as With a fully loaded H3 Q5 GGUF text encoder:
Restoring the upstream delegation (including the final device move) on the same graph produced:
The upstream run did not reproduce the access violation on this current stack. That is only a current-stack non-reproduction; it does not invalidate the older Windows/PyTorch crash reports. The important result for this PR is that resetting This suggests the upstream fix needs to distinguish the |
… move Skipping the base .to(device_to) walk kept a fully loaded model resident on CUDA after a real detach. unpatch_model zeroed the model_loaded_weight_memory and model_offload_buffer_memory counters while the quantized weights never left the device, so the next model saw a full GPU and got pushed into CPU offload. Caught by @woodenriver05 on the PR thread with a 19GB text encoder. Instead of mirroring base's unpatch_weights block minus the device walk, pass device_to=None so base skips only that walk, then move exactly the tensors that report a different device and reassign them through set_attr_param. Host-resident tensors are never touched, which keeps the mmap guarantee from city96#444: mapped memory is only ever read by a cross-device copy, never written in place and never walked by nn.Module.to. device_to=None to super is the same idea as 6dbb4ba, which 717a0e1 reverted for breaking VRAM estimation. The explicit move is the missing half that keeps the accounting correct. Measured on ComfyUI v0.33.1, torch 2.11.0+cu130, Windows 11, RTX 4090, with Wan2.2-I2V-A14B Q5_K_M (10427MB): full load then detach now frees CUDA back to 0MB (stayed at 10512MB before), free_memory after loading a second model frees both (stayed at 20475MB before), and LoRA strength changes 0.25 to 0.75 still reattach patches correctly, matching main byte for byte on the patch state. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Confirmed on my end, and on a full detach it's worse than the partial offload you saw: ComfyUI v0.33.1, torch 2.11.0+cu130, Windows 11, RTX 4090, Wan2.2-I2V-A14B Q5_K_M (10427MB). With 32414a1 a full load then detach left cuda_alloc at 10512MB with both counters zeroed, and free_memory after loading a second copy left 20475MB stranded. Stock frees both to 0. b543b78 drops the mirrored block. The override now passes device_to=None to super so base runs its whole unpatch block and skips only the nn.Module.to walk, then explicitly moves the tensors that report a different device and reassigns them with set_attr_param. Host-resident tensors are never touched, so mmap-backed memory only ever gets read. That covers your partially_load vs detach split by device residency instead of call site: on a real unload the moved set is exactly what has to leave the GPU, on the hot path nothing extra happens. It's also the device_to=None idea from 6dbb4ba that 717a0e1 reverted for breaking VRAM estimation, with the explicit move covering the accounting half. Re-ran the matrix on b543b78: detach frees 10512MB to 0, double-load free_memory 20475MB to 0, LoRA strength 0.25 to 0.75 to detach matches main's patch state exactly, and the fresh, lowvram and cache-flush unpatch paths all complete clean. On the original crash, same result as you: no repro on this stack across five scenario shapes. Probing the mechanism directly, .to() over a still-mmap'd tensor is fine here, but a single in-place write into one is an instant access violation. torch.from_numpy on gguf's read-only mapping gives a tensor torch treats as writable, which is the numpy warning loader.py suppresses. So on current stacks the part that matters is that unpatch never writes mapped memory in place, and this shape preserves that. |
3 participants
Resolves #444.
GGUFModelPatcher.unpatch_modelforwards to baseModelPatcher.unpatch_model, which finishes itsunpatch_weights=Trueblock withself.model.to(device_to). That walk usesnn.Module.toand iterates every parameter, ignoringcomfy_cast_weights = TrueonGGMLLayer. The walk hits the still-mmap'd quantized tensors and faults withEXCEPTION_ACCESS_VIOLATIONon Windows. Repro: any workflow that triggers a cached-patcher flush during sampling (e.g. aUnetLoaderGGUFwired into a multi-output node).The cleanest fix is upstream, since the same conflict will hit any custom op that uses
comfy_cast_weightswith tensors that cannot survivenn.Module.to. Tracked at Comfy-Org/ComfyUI#14142.In the meantime this override inlines base's
unpatch_weights=Trueblock and skips just that one line. The block has to run in full becausepartially_loadreadsmodel_loaded_weight_memoryimmediately afterunpatch_modeland short-circuitsself.load()if it is nonzero - so a partial mirror (the previous attempt at this PR) left.patchescleared without re-attaching anything on LoRA strength changes (caught by @romybaby on this thread). The lines this override now runs map one-for-one to comfy'smodel_patcher.pybody:eject_modelunpatch_hooks,unpin_all_weightsmove_weight_functions/wipe_lowvram_weightcurrent_weight_patches_uuid = None,backup.clear()model_loaded_weight_memory = 0,model_offload_buffer_memory = 0comfy_patched_weightsdeletion per modulesuper().unpatch_model(unpatch_weights=False)foreject_model(idempotent) andobject_patches_backuprestorationDrift risk noted: if upstream adds something to that block, this override needs a sync. Easier to spot since the structure now matches line for line.
Credit @stonerabit for the original repro and diagnosis on #444, and @romybaby for catching the LoRA strength-change regression in the previous iteration.